#include "KDTree.h"
#include <vector>
#include <algorithm>

using namespace std;

struct column_sort
{
    int axis;

    column_sort(int v) : axis(v) {}

    bool operator() (const KDRow & a, const KDRow & b)
    {
        float av = (float)(a.rowData.at<float>(axis));
        float bv = (float)(b.rowData.at<float>(axis));

        return ( av < bv );
    }
};

int KDNode::node_counter = 0;

void delete_kd_tree(KDNode * node)
{
    if ( node )
    {
        delete_kd_tree(node->left_node);
        delete_kd_tree(node->right_node);

        delete node;
    }
}

int getAxisWithMostVariance(const vector<KDRow> & targets)
{

    assert(targets.size() != 0);
    vector<KDRow> rows = targets;
    int num_dim = rows[0].rowData.cols;
    vector<float> means(num_dim);
    for(int d = 0 ; d < num_dim ; d++)
    {
        float sum = 0.0;
        for(int i = 0 ; i < rows.size() ; i++)
        {
            sum += rows[i].rowData.at<float>(d);
        }
        means[d] = sum/rows.size();

    }

    vector<float> variances(num_dim);
    for(int d = 0 ; d < num_dim ; d++)
    {
        float sum = 0.0;
        for(int i = 0 ; i < rows.size() ; i++)
        {
            sum += pow(rows[i].rowData.at<float>(d)-means[d],2);
        }
        variances[d] = sum/rows.size();
    }

    float max_var = -1;
    int max_dim = 0;
    for(int d = 0; d < variances.size() ; d++)
    {
        if(variances[d] > max_var)
        {
            max_dim = d;
            max_var = variances[d];
        }
    }
    return max_dim;

}

KDNode * create_kd_tree(const vector<KDRow> & targets, int depth)
{

    vector<KDRow> rows = targets;

    if ( targets.empty() )
    {

        return NULL;
    }
    else
    {

        if(rows.size() == 1)
        {
            KDNode * pNode = new KDNode();
            pNode->is_leaf = true;
            pNode->originalIdx = rows[0].originalIdx;
            return pNode;
        }

        int axis = getAxisWithMostVariance(targets);
        // cout << axis << endl;
        KDNode * pNode = new KDNode();
        pNode->axis = axis;

        column_sort sorter(axis);

        std::sort(rows.begin(), rows.end(), sorter);

        int median = (int)(rows.size()/2);

        pNode->value = rows[median].rowData.at<float>(axis);
        pNode->originalIdx = rows[median].originalIdx;

        vector<KDRow> left_rows = vector<KDRow>(rows.begin(), rows.begin() + median);
        vector<KDRow> right_rows = vector<KDRow>(rows.begin() + median+1, rows.end());

        pNode->left_node = create_kd_tree(left_rows, depth+1);
        pNode->right_node = create_kd_tree(right_rows, depth+1);

        /*  if ( !pNode->left_node && !pNode->right_node)
        {
            pNode->originalIdx = rows[0].originalIdx;
        }*/

        return pNode;
    }
}

KDTree::KDTree(const cv::Mat & desc)
{
    std::vector<KDRow> targets;



    for(int i=0; i < *(desc.size.p); i++)
    {
        KDRow e;
        e.originalIdx = i;
        e.rowData = desc.row(i);

        targets.push_back(e);
    }

    root_node = create_kd_tree(targets, 0);
}

KDTree::~KDTree()
{
    delete root_node;
}

KDNode * search_kd_tree(KDNode * pNode, int depth,  const cv::Mat & original, const cv::Mat & target, float maxDistance, float & distance)
{
    if ( pNode->is_leaf)
    {
        // L2 distance measure
        distance = 0;

        //cv::Mat diff = original.row(pNode->originalIdx) - target;

        float diff = 0;

        for(int i=0; i < target.cols; i++)
        {
            float a = original.row(pNode->originalIdx).at<float>(i);
            float b = target.at<float>(i);

            diff += (a - b) * (a - b);
        }

        distance = sqrt(diff);

       // cout << "leaf dist: " << distance << endl;
        //if ( distance < maxDistance )*/
        return pNode;
        //else
        //    return NULL;
    }

    float node_distance = 0;
    float diff = 0;

    for(int i=0; i < target.cols; i++)
    {
        float a = original.row(pNode->originalIdx).at<float>(i);
        float b = target.at<float>(i);

        diff += (a - b) * (a - b);
    }

    node_distance = sqrt(diff);
  //  cout << "node_dist: " << node_distance << endl;

    int axis = pNode->axis;

    float kdtree_value = pNode->value;
    float search_value = target.at<float>(axis);

    if ( search_value < kdtree_value && pNode->left_node)
    {
        KDNode * current_best = search_kd_tree(pNode->left_node, depth+1, original, target, maxDistance, distance);
        float dist = fabs(original.row(pNode->originalIdx).at<float>(axis)-target.at<float>(axis));
        if(distance < dist)
        {
            //cout << "1 dist: " << distance << endl;
            return current_best;
        }
        else if(pNode->right_node)
        {
            float other_distance= 0;
            KDNode * other_best = search_kd_tree(pNode->right_node, depth+1, original, target, maxDistance, other_distance);
            if(other_distance < distance)
            {

                if(distance < node_distance)
                {
                  //  cout << "2 dist: " << distance << endl;
                    distance = other_distance;
                    return other_best;
                }


            }
        }


        if(distance < node_distance)
        {
            return current_best;
        }

    }
    else if(pNode->right_node)
    {
        KDNode * current_best = search_kd_tree(pNode->right_node, depth+1, original, target, maxDistance, distance);
        float dist = fabs(original.row(pNode->originalIdx).at<float>(axis)-target.at<float>(axis));
        if(distance < dist)
        {
            //cout << "3 dist: " << distance << endl;
            return current_best;
        }
        else if(pNode->left_node)
        {
            float other_distance= 0;
            KDNode * other_best = search_kd_tree(pNode->left_node, depth+1, original, target, maxDistance, other_distance);
            if(other_distance < distance)
            {
                distance = other_distance;
                if(distance < node_distance)
                {
                    //cout << "4 dist: " << distance << endl;
                    return other_best;
                }
            }
        }

        if(distance < node_distance)
        {
            return current_best;
        }
    }

    distance = node_distance;

    //cout << "end dist: " << distance << endl;
    return pNode;

}

void KDTree::match(   const cv::Mat & original,const cv::Mat& target, std::vector<cv::DMatch>& match_result, float maxDistance)
{
    match_result.clear();
   // cout << "original: " << original.type() <<" CV_64F: " << CV_64F << " CV_32F: " << CV_32F << endl;

    for(int i=0; i< *(target.size.p); i++)
    {
        float distance;
        KDNode * pSearchResult = search_kd_tree(root_node, 0,original, target.row(i), maxDistance, distance);

        if ( pSearchResult)
        {
            //if(distance < maxDistance)
            cout << distance << "\n";
            {
                cv::DMatch match;
                match.trainIdx = pSearchResult->originalIdx;
                match.queryIdx = i;
                match.distance = distance;

                match_result.push_back(match);
            }
        }
    }
    cout << flush;
}
